package algorithm_demo.demo04;

import java.util.*;

/**
 * 并查集
 *
 * @author Api
 * @date 2023/2/18 12:06
 */
public class Code07_UnionFind {
    /*并查集
     * 1. 有若干个样本a、b、c、d...类型假设是V
     * 2. 在并查集中一开始认为每个样本都在单独的集合里
     * 3. 用户可以在任何时候调用如下两个方法：
     *     boolean isSameSet(V x, V y)：查询样本x和样本y是否属于一个集合
     *     void union(V x, V y):把x和y各自所在集合的所有样本合并成一个集合
     * 4. isSameSet和union方法的代价越低越好
     * */

    public static class Node<V> {
        public V value;

        public Node(V v) {
            this.value = v;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }
            Node<?> node = (Node<?>) o;
            return Objects.equals(value, node.value);
        }

        @Override
        public int hashCode() {
            return Objects.hash(value);
        }
    }


    public static class UnionFind<V> {
        //V->节点
        public HashMap<V, Node<V>> nodes;
        public HashMap<Node<V>, Node<V>> parents;
        //只有一个点，它是代表点，才有记录
        public HashMap<Node<V>, Integer> sizeMap;

        public UnionFind(Collection<V> values) {
            nodes = new HashMap<>();
            parents = new HashMap<>();
            sizeMap = new HashMap<>();
            for (V cur : values) {
                Node<V> node = new Node<>(cur);
                nodes.put(cur, node);//一一对应，不会更改
                parents.put(node, node);
                sizeMap.put(node, 1);
            }
        }

        //从点cur开始，一直往上找，找到不能再往上的代表点(头节点)，返回
        public Node<V> findFather(Node<V> cur) {
            Stack<Node<V>> path = new Stack<>();
            while (cur != parents.get(cur)) {
                path.push(cur);
                cur = parents.get(cur);
            }
            //cur肯定是头结点
            //扁平化过程，cur表示代表点，将栈弹出的每一个元素
            while (!path.isEmpty()) {
                parents.put(path.pop(), cur);
            }
            return cur;
        }

        public boolean isSameSet(V a, V b) {
            //a和b是否进行过注册
            if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
                return false;
            }
            return findFather(nodes.get(a)) == findFather(nodes.get(b));
        }

        /**
         * 前提是需要把样本都放入到nodes的HashMap集合中
         *
         * @param a
         * @param b
         */
        public void union(V a, V b) {
            //未注册直接返回
            if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
                return;
            }
            Node<V> aHead = findFather(nodes.get(a));
            Node<V> bHead = findFather(nodes.get(b));
            //校验aHead和bHead内存地址相同时，直接跳过
            if (aHead != bHead) {
                int aSetSize = sizeMap.get(aHead);
                int bSetSize = sizeMap.get(bHead);
                if (aSetSize >= bSetSize) {//如果aSetSize集合>=bSetSize
                    parents.put(bHead, aHead);//小挂大，需要把小集合设置为大集合中
                    sizeMap.put(aHead, aSetSize + bSetSize);
                    sizeMap.remove(bHead);
                } else {//如果aSetSize集合<bSetSize
                    parents.put(aHead, bHead);
                    sizeMap.put(bHead, aSetSize + bSetSize);
                    sizeMap.remove(aHead);
                }
            }
        }

        public int getSetNums() {
            return sizeMap.size();
        }
    }
}
